from pyspark.sql import SparkSession
from pyspark.sql.functions import col, countDistinct

# 初始化Spark会话
spark = SparkSession.builder \
    .appName("ShopCustomerCount") \
    .getOrCreate()

def calculate_shop_customer_count(
        visit_data_path: str,
        short_video_interaction_path: str,
        cart_payment_path: str,
        start_date: str,
        end_date: str
) -> int:
    """
    计算指定日期范围内的店铺客户数

    参数:
    - visit_data_path: 进店访问数据路径
    - short_video_interaction_path: 短视频互动数据路径
    - cart_payment_path: 购物车支付数据路径
    - start_date: 开始日期 (格式: 'YYYY-MM-DD')
    - end_date: 结束日期 (格式: 'YYYY-MM-DD')

    返回:
    - 店铺客户数
    """
    # 读取进店访问数据（仅无线端）
    visit_df = spark.read.csv(visit_data_path,header=True, inferSchema=True) \
        .filter(col("visit_time").between(start_date, end_date)) \
        .filter(col("terminal_type") == "wireless") \
        .select("user_id") \
        .distinct()

    # 读取短视频互动数据（观看<3秒且有点赞行为的无线端用户）
    video_df = spark.read.csv(short_video_interaction_path,header=True, inferSchema=True) \
        .filter(col("interaction_time").between(start_date, end_date)) \
        .filter(col("watch_duration") < 3) \
        .filter(col("action_type") == "like") \
        .filter(col("terminal_type") == "wireless") \
        .select("user_id") \
        .distinct()

    # 读取购物车支付数据（未进店直接支付的无线端用户）
    cart_df = spark.read.csv(cart_payment_path,header=True, inferSchema=True) \
        .filter(col("payment_time").between(start_date, end_date)) \
        .filter(col("is_direct_payment") == True) \
        .filter(col("terminal_type") == "wireless") \
        .select("user_id") \
        .distinct()

    # 合并三类用户并去重
    all_users_df = visit_df.union(video_df).union(cart_df).distinct()

    # 计算店铺客户数
    customer_count = all_users_df.count()

    return customer_count

# 示例调用
if __name__ == "__main__":
    customer_count = calculate_shop_customer_count(
        visit_data_path="mock_data/visit_data.csv",
        short_video_interaction_path="mock_data/video_interaction_data.csv",
        cart_payment_path="mock_data/cart_payment_data.csv",
        start_date="2025-07-01",
        end_date="2025-07-10"
    )

    print(f"店铺客户数 (2025-07-01至2025-07-10): {customer_count}")